package cas.ihep.spark.util;

import java.io.Serializable;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;

public class Counter<P> implements Serializable,Iterable<Map.Entry<P,Integer>> {

    private Map<P,int[]> cnt;

    public Counter(){
        cnt=new HashMap<>();
    }

    public int count(P key){
        int[] valueWrapper=cnt.get(key);
        if(valueWrapper==null){
            cnt.put(key,new int[]{1});
            return 1;
        }else{
            return ++valueWrapper[0];
        }
    }

    public int get(P key){
        int[] valueWrapper=cnt.get(key);
        return valueWrapper!=null?valueWrapper[0]:0;
    }

    public Collection<Map.Entry<P,Integer>> values(){
        HashMap<P,Integer> ret=new HashMap<>();
        for(Map.Entry<P,int[]> entry:cnt.entrySet()){
            ret.put(entry.getKey(),entry.getValue()[0]);
        }
        return ret.entrySet();
    }

    @SuppressWarnings("NullableProblems")
    @Override
    public Iterator<Map.Entry<P, Integer>> iterator() {
        return values().iterator();
    }

    public void forEach(BiConsumer<P> consumer){
        for(Map.Entry<P,int[]> entry:cnt.entrySet()){
            consumer.call(entry.getKey(),entry.getValue()[0]);
        }
    }

    public interface BiConsumer<P>{
        void call(P key,int value);
    }

    public int size(){
        return cnt.size();
    }
}
